#include "S5ActionClient.h"
#include <KSocket.h>
#include "S5ActionBind.h"
#include "S5ActionConnect.h"
#include "S5ActionServer.h"
#include "S5ActionUdp.h"
#include "S5Debug.h"
#include "S5Def.h"
#include "S5Server.h"

S5ActionClient::S5ActionClient(const KIAddress &addr, S5ActionServer *server) :
    KTimer(server->server()->queue()),
    m_abortOnSent(false), m_addr(addr), m_server(server),
    m_proxy_receiver(NULL), m_proxy_sender(NULL),
    m_storage(server->server()->pool()),
    m_state(S5ActionClient::TIMEOUT), m_length(1)
{
    start(S5_TIME_OUT);
}

S5ActionClient::~S5ActionClient()
{
    assert(!isAlive());
}

std::string S5ActionClient::domain() const
{
    if (m_addr_type == S5_ATYP_IPV4) {
        return KIpAddress(0, m_dest_addr).ip();
    } else {
        return m_domain;
    }
}

void S5ActionClient::relay(const void *data, size_t length, const KIAddress *dest)
{
    S5DebugHex(data, length, "Tcp send to client", &m_addr);
    server()->send(socket(), data, length);
}

S5ActionClient::Result S5ActionClient::process(void *data, size_t length)
{
    if (m_state == S5ActionClient::TCP_RELAY) {
        if (S5Action *proxy = m_proxy_sender) {
            S5DebugHex(data, length, "Tcp receive from client", &m_addr);
            proxy->relay(data, length, NULL);
        }
        return S5ActionClient::Idle;
    } else if (m_state == S5ActionClient::UDP_RELAY) {
        KDebugMsg("Warning: ignore tcp data while udp relay");
    } else if (m_storage.write(data, length)) {
        reset(S5_TIME_OUT);
        while (m_storage.bytesAvailable() >= m_length) {
            switch (m_state) {
            case S5ActionClient::TIMEOUT:
                m_state = S5ActionClient::VER_1;
            case S5ActionClient::VER_1:
                {
                    if (m_storage.peekChar() == S5_VERSION_PROTOCOL) {
                        m_state = S5ActionClient::NMETHODS;
                        m_length = 1;
                        m_storage.ignore(1);
                    } else {
                        return S5ActionClient::Abort;
                    }
                } break;
            case S5ActionClient::NMETHODS:
                {
                    m_state = S5ActionClient::METHODS;
                    m_length = m_storage.peekChar();
                    m_storage.ignore(1);
                } break;
            case S5ActionClient::METHODS:
                {
                    MethodAck ack;
                    if (m_length > 0) {
                        if (server()->requireAuthentication()) {
                            do {
                                // Username/Password Authentication
                                if (m_storage.peekChar() == S5_METHOD_USERNAME_PASSWORD) {
                                    ack.method = S5_METHOD_USERNAME_PASSWORD;
                                    m_storage.ignore(m_length);
                                    break;
                                } else {
                                    m_storage.ignore(1);
                                }
                            } while (--m_length);
                        } else {
                            ack.method = S5_METHOD_NO_AUTHENTICATION;
                            m_storage.ignore(m_length);
                        }
                    }
                    switch (ack.method) {
                    case S5_METHOD_NO_AUTHENTICATION:
                        {
                            m_state = S5ActionClient::VER_3;
                            m_length = 1;
                            server()->sendAck(socket(), &ack, sizeof ack);
                        } break;
                    case S5_METHOD_USERNAME_PASSWORD:
                        {
                            m_state = S5ActionClient::VER_2;
                            m_length = 1;
                            server()->sendAck(socket(), &ack, sizeof ack);
                        } break;
                    default:
                        {
                            m_abortOnSent = true;
                            server()->sendAck(socket(), &ack, sizeof ack);
                            return S5ActionClient::Idle;
                        } break;
                    }
                } break;
            case S5ActionClient::VER_2:
                {
                    if (m_storage.peekChar() == S5_VERSION_AUTHENTICATION) {
                        m_state = S5ActionClient::ULEN;
                        m_length = 1;
                        m_storage.ignore(1);
                    } else {
                        return S5ActionClient::Abort;
                    }
                } break;
            case S5ActionClient::ULEN:
                {
                    m_state = S5ActionClient::UNAME;
                    m_length = m_storage.peekChar();
                    m_storage.ignore(1);
                } break;
            case S5ActionClient::UNAME:
                {
                    if (m_length > 0) {
                        do {
                            m_username += m_storage.peekChar();
                            m_storage.ignore(1);
                        } while (--m_length);
                    }
                    m_state = S5ActionClient::PLEN;
                    m_length = 1;
                } break;
            case S5ActionClient::PLEN:
                {
                    m_state = S5ActionClient::PASSWD;
                    m_length = m_storage.peekChar();
                    m_storage.ignore(1);
                } break;
            case S5ActionClient::PASSWD:
                {
                    if (m_length > 0) {
                        do {
                            m_password += m_storage.peekChar();
                            m_storage.ignore(1);
                        } while (--m_length);
                    }
                    AuthenticationAck ack;
                    if (server()->verify(m_username, m_password)) {
                        ack.status = S5_STATUS_SUCCESS;
                        m_state = S5ActionClient::VER_3;
                        m_length = 1;
                        server()->sendAck(socket(), &ack, sizeof ack);
                    } else {
                        m_abortOnSent = true;
                        server()->sendAck(socket(), &ack, sizeof ack);
                        return S5ActionClient::Idle;
                    }
                } break;
            case S5ActionClient::VER_3:
                {
                    if (m_storage.peekChar() == S5_VERSION_PROTOCOL) {
                        m_state = S5ActionClient::CMD;
                        m_length = 1;
                        m_storage.ignore(1);
                    } else {
                        return S5ActionClient::Abort;
                    }
                } break;
            case S5ActionClient::CMD:
                {
                    switch ((m_command = m_storage.peekChar())) {
                    case S5_CMD_CONNECT:
                    case S5_CMD_BIND:
                    case S5_CMD_UDP_ASSOCIATE:
                        {
                            m_state = S5ActionClient::RSV;
                            m_length = 1;
                            m_storage.ignore(1);
                        } break;
                    default:
                        {
                            RequestAck ack;
                            ack.rep = S5_REP_COMMAND_NOT_SUPPORTED;
                            m_abortOnSent = true;
                            server()->sendAck(socket(), &ack, sizeof ack);
                            return S5ActionClient::Idle;
                        } break;
                    }
                } break;
            case S5ActionClient::RSV:
                {
                    if (m_storage.peekChar() == S5_RSV) {
                        m_state = S5ActionClient::ATYP;
                        m_length = 1;
                        m_storage.ignore(1);
                    } else {
                        return S5ActionClient::Abort;
                    }
                } break;
            case S5ActionClient::ATYP:
                {
                    switch ((m_addr_type = m_storage.peekChar())) {
                    case S5_ATYP_IPV4:
                        {
                            m_state = S5ActionClient::DST_ADDR;
                            m_length = 4;
                            m_storage.ignore(1);
                        } break;
                    case S5_ATYP_DOMAINNAME:
                        {
                            m_state = S5ActionClient::DLEN;
                            m_length = 1;
                            m_storage.ignore(1);
                        } break;
                    default:
                        {
                            RequestAck ack;
                            ack.rep = S5_REP_ADDRESS_NOT_SUPPORTED;
                            m_abortOnSent = true;
                            server()->sendAck(socket(), &ack, sizeof ack);
                            return S5ActionClient::Idle;
                        } break;
                    }
                } break;
            case S5ActionClient::DLEN:
                {
                    m_state = S5ActionClient::DST_ADDR;
                    m_length = m_storage.peekChar();
                    m_storage.ignore(1);
                } break;
            case S5ActionClient::DST_ADDR:
                {
                    if (m_addr_type == S5_ATYP_IPV4) {
                        assert(sizeof m_dest_addr == m_length);
                        m_storage.read(&m_dest_addr, m_length);
                    } else {
                        if (m_length > 0) {
                            do {
                                m_domain += m_storage.peekChar();
                                m_storage.ignore(1);
                            } while (--m_length);
                        }
                        KTcpHostAddress addr(m_domain.c_str(), NULL);
                        if (addr.isNull()) {
                            RequestAck ack;
                            ack.rep = S5_REP_HOST_UNREACHABLE;
                            m_abortOnSent = true;
                            server()->sendAck(socket(), &ack, sizeof ack);
                            return S5ActionClient::Idle;
                        } else {
                            m_dest_addr = KIpAddress(addr).netIp();
                        }
                    }
                    m_state = S5ActionClient::DST_PORT;
                    m_length = 2;
                } break;
            case S5ActionClient::DST_PORT:
                {
                    assert(sizeof m_dest_port == m_length);
                    m_storage.read(&m_dest_port, m_length);
                    stop();
                    m_storage.clear();
                    m_storage.free();
                    m_state = S5ActionClient::EVAL;
                    return evaluate();
                } break;
            default:
                assert(false);
                break;
            }
        }
    } else {
        return S5ActionClient::Abort;
    }
    return S5ActionClient::Receive;
}

S5ActionClient::Result S5ActionClient::evaluate()
{
    assert(m_state == S5ActionClient::EVAL);
    if (socket()->isAlive()) {
        server()->onEval(this);
        KIpAddress addr(m_dest_port, m_dest_addr), local(0);
        switch (m_command) {
        case S5_CMD_CONNECT:
            {
                m_proxy_sender = new S5ActionConnect(this);
                server()->addConnector(local, addr, m_proxy_sender);
            } break;
        case S5_CMD_BIND:
            {
                if (KSocket *sock = server()->addAcceptor(local)) {
                    m_proxy_sender = new S5ActionBind(this);
                    server()->setAction(sock, m_proxy_sender);
                    if (sock->getLocalAddress(addr)) {
                        postSuccess(addr, true);
                    } else {
                        postFailure(S5_REP_GENERAL_SERVER_FAILURE, false);
                    }
                } else {
                    postFailure(S5_REP_GENERAL_SERVER_FAILURE);
                }
            } break;
        case S5_CMD_UDP_ASSOCIATE:
            {
                RequestAck ack;
                if (KSocket *receiver = server()->addMessenger(local)) {
                    if (KSocket *sender = server()->addMessenger(local)) {
                        S5ActionUdp *context_recv = new S5ActionUdp(this);
                        S5ActionUdp *context_send = new S5ActionUdp(this);
                        context_recv->associate(context_send);
                        context_recv->setSourceAddress(addr);
                        server()->ensureAlive(receiver);
                        server()->ensureAlive(sender);
                        server()->setAction(receiver, context_recv);
                        server()->setAction(sender, context_send);
                        m_proxy_receiver = context_recv;
                        m_proxy_sender = context_send;
                        {
                            KIpAddress bnd;
                            receiver->getLocalAddress(bnd);
                            ack.rep = S5_REP_SUCCEEDED;
                            ack.bnd_addr = bnd.netIp();
                            ack.bnd_port = bnd.netPort();
                            m_state = S5ActionClient::UDP_RELAY;
                            server()->receive(receiver);
                            server()->receive(sender);
                            server()->sendAck(socket(), &ack, sizeof ack);
                        }
                        return S5ActionClient::Receive;
                    } else {
                        receiver->shutdown();
                        receiver->close();
                    }
                }
                ack.rep = S5_REP_GENERAL_SERVER_FAILURE;
                m_abortOnSent = true;
                server()->sendAck(socket(), &ack, sizeof ack);
            } break;
        default:
            assert(false);
            break;
        }
    }
    return S5ActionClient::Idle;
}

void S5ActionClient::postFailure(uint8_t rep, bool clear)
{
    if (clear) {
        m_proxy_sender = NULL;
    }
    RequestAck ack;
    ack.rep = rep;
    m_abortOnSent = true;
    server()->sendAck(socket(), &ack, sizeof ack);
}

void S5ActionClient::postSuccess(const KIpAddress &addr, bool wait)
{
    RequestAck ack;
    ack.rep = S5_REP_SUCCEEDED;
    ack.bnd_addr = addr.netIp();
    ack.bnd_port = addr.netPort();
    if (wait) {
        server()->sendAck(socket(), &ack, sizeof ack);
    } else {
        m_state = S5ActionClient::TCP_RELAY;
        if (server()->sendAck(socket(), &ack, sizeof ack)) {
            socket()->setNoDelay(false);
            server()->receive(socket());
        }
    }
}

void S5ActionClient::abortImp()
{
    // Close client side
    socket()->shutdown();
    socket()->close();
    // Close proxy side
    if (m_proxy_receiver) {
        m_proxy_receiver->abortOnce();
        m_proxy_receiver = NULL;
    }
    if (m_proxy_sender) {
        m_proxy_sender->abortOnce();
        m_proxy_sender = NULL;
    }
    // Notify server
    server()->onAborted(this);
    server()->ensureAlive(socket(), false);
    // Update server context
    m_server->decreaseClients();
}

void S5ActionClient::onClosing(bool stopped)
{
    stop();
}

void S5ActionClient::onError(KDef::SockError type, DWORD error)
{
    abortOnce();
}

void S5ActionClient::onReceived(KKey *&key, void *data, size_t length, bool more)
{
    S5Action::onReceived(key, data, length, more);
    switch (process(data, length)) {
    case S5ActionClient::Abort:
        abortOnce();
        break;
    case S5ActionClient::Receive:
        server()->receive(socket());
        break;
    }
}

void S5ActionClient::onSent(KKey *&key, void *data, size_t length)
{
    S5Action::onSent(key, data, length);
    if (m_abortOnSent) {
        abortOnce();
    } else if (!key) {
        if (S5Action *proxy = m_proxy_sender) {
            proxy->onPeerRelayed();
        }
    }
}

void S5ActionClient::tick()
{
    // Negotiation time out
    abortOnce();
}